from typing import ClassVar
from pandas import Series, Index
from sklearn.feature_selection import GenericUnivariateSelect, f_regression, mutual_info_regression
from selector._base_selector import _BaseSelector


class Filter(_BaseSelector):
    """
    Filter method of feature selection, subclass of _BaseFeatureSelector.
    Parameters:
        same with father class
    """

    def __init__(self, x, y, n_selected_features):
        super(Filter, self).__init__(x, y, n_selected_features)

    def select_features(self, score_func: ClassVar, mode: str) -> Index:
        """
        Parameters:
            score_func: str,
                for regression: {'f_regression', 'mutual_info_regression'}
                for classification: {'chi2', 'f_classif', 'mutual_info_classif'}
            mode: str, {'percentile', 'k_best', 'fpr', 'fdr', 'fwe'}, 'percentile' not used here
            (param: float or int, depending on the feature selection mode)
        Return:
            Index, names of sorted selected features
        """
        param = self.n_selected_features if mode == 'k_best' else 0.05
        selector = GenericUnivariateSelect(score_func=eval(score_func), mode=mode, param=param)
        selector.fit(self.x, self.y)
        return super(Filter, self).sort_features(selector.scores_)


class CorrelationFilter(_BaseSelector):
    """
    Filter method by pandas correlation method of feature selection, subclass of _BaseFeatureSelector.
    Parameters:
        same with father class
    """

    def __init__(self, x, y, n_selected_features):
        super(CorrelationFilter, self).__init__(x, y, n_selected_features)

    def select_features(self, method: str) -> Index:
        """
        Parameters:
            method: str, {'pearson', 'kendall', 'spearman'}
        Return:
            Index, names of sorted selected features
        """
        scores = Series(index=self.x.columns, dtype='float')
        for col in self.x.columns:
            scores[col] = self.y.corr(self.x.loc[:, col], method=method)
        return super(CorrelationFilter, self).sort_features(scores)
